import datetime
import json
import queue
import socket
import threading
import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox, filedialog
from typing import Dict
import ssl
import os
import tempfile
from OpenSSL import crypto
import time
import re


class HTTPMessage:
    def __init__(self, raw_data: bytes):
        self.raw_data = raw_data
        self.method = ""
        self.path = ""
        self.status_code = ""
        self.status_text = ""
        self.headers: Dict[str, str] = {}
        self.body = ""
        self.is_request = True
        self.parse_message()

    def parse_message(self):
        try:
            # 将字节转换为字符串，保留二进制数据
            try:
                # 先尝试 UTF-8 解码
                message_str = self.raw_data.decode('utf-8', errors='ignore')
            except UnicodeDecodeError:
                # 如果失败，尝试其他编码
                try:
                    message_str = self.raw_data.decode('gbk', errors='ignore')
                except UnicodeDecodeError:
                    message_str = self.raw_data.decode('latin1', errors='ignore')

            # 分离请求行/状态行、头部和主体
            if '\r\n\r\n' in message_str:
                headers_part, self.body = message_str.split('\r\n\r\n', 1)
            else:
                headers_part = message_str
                self.body = ''

            lines = headers_part.split('\r\n')
            if not lines:
                return

            # 解析第一行
            first_line = lines[0].split(' ')
            if len(first_line) >= 3:
                if first_line[0].startswith('HTTP/'):
                    # 这是一个响应
                    self.is_request = False
                    self.status_code = first_line[1]
                    self.status_text = ' '.join(first_line[2:])
                else:
                    # 这是一个请求
                    self.method = first_line[0]
                    self.path = first_line[1]

            # 解析头部
            for line in lines[1:]:
                if ': ' in line:
                    key, value = line.split(': ', 1)
                    self.headers[key] = value

            # 处理Content-Type，尝试解析不同格式的响应体
            content_type = self.headers.get('Content-Type', '').lower()
            if content_type and self.body:
                if 'json' in content_type:
                    try:
                        # 格式化JSON数据
                        json_data = json.loads(self.body)
                        self.body = json.dumps(json_data, indent=2, ensure_ascii=False)
                    except:
                        pass
                elif 'xml' in content_type:
                    try:
                        # 尝试格式化XML
                        from xml.dom.ext import PrettyPrint
                        from xml.dom.ext.reader import Sax2
                        from io import StringIO
                        reader = Sax2.Reader()
                        xml_doc = reader.fromString(self.body)
                        stream = StringIO()
                        PrettyPrint(xml_doc, stream)
                        self.body = stream.getvalue()
                    except:
                        pass

        except Exception as e:
            print(f"解析HTTP消息时出错: {e}")

    def __str__(self):
        if self.is_request:
            return f"{self.method} {self.path}"
        else:
            return f"{self.status_code} {self.status_text}"

class CertificateAuthority:
    def __init__(self):
        self.ca_cert_path = os.path.join(tempfile.gettempdir(), "proxy_ca.crt")
        self.ca_key_path = os.path.join(tempfile.gettempdir(), "proxy_ca.key")
        self.create_ca_cert()

    def create_ca_cert(self):
        """创建CA证书"""
        if os.path.exists(self.ca_cert_path) and os.path.exists(self.ca_key_path):
            return

        # 生成密钥
        key = crypto.PKey()
        key.generate_key(crypto.TYPE_RSA, 2048)

        # 创建CA证书
        cert = crypto.X509()
        cert.get_subject().CN = "Proxy CA"
        cert.set_serial_number(1000)
        cert.gmtime_adj_notBefore(0)
        cert.gmtime_adj_notAfter(315360000)  # 10年有效期
        cert.set_issuer(cert.get_subject())
        cert.set_pubkey(key)
        cert.sign(key, 'sha256')

        # 保存证书和私钥
        with open(self.ca_cert_path, "wb") as f:
            f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert))
        with open(self.ca_key_path, "wb") as f:
            f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, key))

    def create_cert(self, hostname):
        """为指定域名创建证书"""
        # 读取CA证书和私钥
        with open(self.ca_cert_path, 'rb') as f:
            ca_cert = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
        with open(self.ca_key_path, 'rb') as f:
            ca_key = crypto.load_privatekey(crypto.FILETYPE_PEM, f.read())

        # 生成密钥
        key = crypto.PKey()
        key.generate_key(crypto.TYPE_RSA, 2048)

        # 创建证书
        cert = crypto.X509()
        cert.get_subject().CN = hostname
        cert.set_serial_number(1000)
        cert.gmtime_adj_notBefore(0)
        cert.gmtime_adj_notAfter(31536000)  # 1年有效期
        cert.set_issuer(ca_cert.get_subject())
        cert.set_pubkey(key)
        cert.sign(ca_key, 'sha256')

        return cert, key

class SecurityManager:
    def __init__(self):
        # IP黑名单
        self.blacklist = set()
        # IP请求计数器
        self.request_counters = {}
        # 上次请求时间记录
        self.last_request_time = {}
        # 配置参数
        self.config = {
            'max_requests_per_minute': 60,  # 每分钟最大请求数
            'min_request_interval': 0.1,    # 最小请求间隔（秒）
            'max_payload_size': 1024 * 1024,  # 最大请求体积 1MB
            'max_uri_length': 2048,         # 最大URI长度
            'max_header_count': 50,         # 最大请求头数量
            'max_header_size': 8192,        # 最大请求头大小
        }
        # 攻击特征
        self.attack_patterns = [
            # SQL注入特征
            r"(?i)(SELECT|INSERT|UPDATE|DELETE|DROP|UNION|OR 1=1)",
            # XSS特征
            r"(?i)(<script|javascript:|on\w+\s*=)",
            # 命令注入特征
            r"(?i)([;&|`]|\$\(|\|\||&&)",
            # 路径遍历特征
            r"(?i)(\.\.\/|\.\.\\)",
            # 文件包含特征
            r"(?i)(include\(|require\(|file:\/\/)",
        ]
        # 清理计数器的定时器
        self._start_cleanup_timer()

    def _start_cleanup_timer(self):
        """启动定时清理任务"""
        def cleanup():
            current_time = time.time()
            # 清理超过1分钟的计数器
            self.request_counters = {
                ip: (count, timestamp) 
                for ip, (count, timestamp) in self.request_counters.items()
                if current_time - timestamp < 60
            }
            # 清理超过1分钟的时间记录
            self.last_request_time = {
                ip: timestamp 
                for ip, timestamp in self.last_request_time.items()
                if current_time - timestamp < 60
            }
            # 每分钟执行一次清理
            threading.Timer(60, cleanup).start()
        
        cleanup()

    def check_request(self, client_address: str, request_data: bytes) -> tuple[bool, str]:
        """检查请求是否安全
        返回: (是否通过, 拒绝原因)
        """
        try:
            # 检查IP是否在黑名单中
            if client_address in self.blacklist:
                return False, "IP已被封禁"

            # 检查请求频率
            current_time = time.time()
            if client_address in self.last_request_time:
                time_diff = current_time - self.last_request_time[client_address]
                if time_diff < self.config['min_request_interval']:
                    return False, "请求过于频繁"

            # 更新最后请求时间
            self.last_request_time[client_address] = current_time

            # 检查请求计数
            if client_address in self.request_counters:
                count, timestamp = self.request_counters[client_address]
                if current_time - timestamp < 60:  # 1分钟内
                    if count >= self.config['max_requests_per_minute']:
                        self.blacklist.add(client_address)  # 加入黑名单
                        return False, "请求次数超限"
                    self.request_counters[client_address] = (count + 1, timestamp)
                else:
                    self.request_counters[client_address] = (1, current_time)
            else:
                self.request_counters[client_address] = (1, current_time)

            # 检查请求大小
            if len(request_data) > self.config['max_payload_size']:
                return False, "请求体积过大"

            # 解析请求进行更详细的检查
            try:
                request_str = request_data.decode('utf-8', errors='ignore')
                request_lines = request_str.split('\r\n')
                
                # 检查请求行
                if request_lines:
                    first_line = request_lines[0]
                    # 检查URI长度
                    if len(first_line) > self.config['max_uri_length']:
                        return False, "URI长度超限"
                    
                    # 检查攻击特征
                    for pattern in self.attack_patterns:
                        if re.search(pattern, first_line):
                            self.blacklist.add(client_address)  # 发现攻击特征，加入黑名单
                            return False, "检测到潜在攻击"

                # 检查请求头
                headers_count = 0
                total_headers_size = 0
                for line in request_lines[1:]:
                    if not line:
                        break
                    headers_count += 1
                    total_headers_size += len(line)
                    
                    # 检查请求头数量
                    if headers_count > self.config['max_header_count']:
                        return False, "请求头数量超限"
                    
                    # 检查请求头大小
                    if total_headers_size > self.config['max_header_size']:
                        return False, "请求头大小超限"
                    
                    # 检查攻击特征
                    for pattern in self.attack_patterns:
                        if re.search(pattern, line):
                            self.blacklist.add(client_address)
                            return False, "检测到潜在攻击"

            except Exception as e:
                print(f"解析请求时出错: {e}")
                return False, "请求格式错误"

            return True, "通过"

        except Exception as e:
            print(f"安全检查时出错: {e}")
            return False, "安全检查失败"

    def remove_from_blacklist(self, ip: str):
        """从黑名单中移除IP"""
        if ip in self.blacklist:
            self.blacklist.remove(ip)
            if ip in self.request_counters:
                del self.request_counters[ip]
            if ip in self.last_request_time:
                del self.last_request_time[ip]

class NetworkSniffer:
    def __init__(self, host='127.0.0.1', port=8888, callback=None, listen_all=False):
        self.host = host
        self.port = port
        self.callback = callback
        self.running = False
        self.listen_all = listen_all
        self.proxy_socket = None
        self.connections = {}
        
        # 安全管理器
        self.security_manager = SecurityManager()
        
        # 断点控制
        self.break_on_request = False
        self.break_on_response = False
        self.breakpoint_event = threading.Event()
        self.breakpoint_data = None
        self.continue_with_modified = False

        # 初始化证书颁发机构
        self.ca = CertificateAuthority()
        
        # 证书缓存
        self.cert_cache = {}

    def is_port_in_use(self, port):
        """检查端口是否被占用"""
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind(('127.0.0.1', port))
                return False
            except socket.error:
                return True

    def force_release_port(self, port):
        """强制释放端口"""
        try:
            # Windows系统
            if os.name == 'nt':
                os.system(f'netstat -ano | findstr :{port} > port_info.txt')
                with open('port_info.txt', 'r') as f:
                    for line in f:
                        if f':{port}' in line:
                            pid = line.strip().split()[-1]
                            os.system(f'taskkill /F /PID {pid}')
                os.remove('port_info.txt')
            # Linux/Unix系统
            else:
                os.system(f'lsof -i :{port} | grep LISTEN | awk \'{{print $2}}\' | xargs kill -9')
            return True
        except Exception as e:
            print(f"释放端口失败: {e}")
            return False

    def create_ssl_context(self, hostname):
        """创建SSL上下文"""
        if hostname not in self.cert_cache:
            cert, key = self.ca.create_cert(hostname)
            self.cert_cache[hostname] = (cert, key)
        else:
            cert, key = self.cert_cache[hostname]

        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        cert_file = os.path.join(tempfile.gettempdir(), f"{hostname}.crt")
        key_file = os.path.join(tempfile.gettempdir(), f"{hostname}.key")

        # 保存证书和私钥到临时文件
        with open(cert_file, "wb") as f:
            f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert))
        with open(key_file, "wb") as f:
            f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, key))

        context.load_cert_chain(cert_file, key_file)
        
        # 清理临时文件
        os.unlink(cert_file)
        os.unlink(key_file)
        
        return context

    def set_breakpoint(self, request=None, response=None):
        """设置断点状态"""
        if request is not None:
            self.break_on_request = request
        if response is not None:
            self.break_on_response = response

    def wait_for_breakpoint(self, data):
        """等待断点处理"""
        self.breakpoint_data = data
        self.breakpoint_event.clear()
        self.breakpoint_event.wait()  # 等待继续信号
        modified_data = self.breakpoint_data if self.continue_with_modified else data
        self.breakpoint_data = None
        self.continue_with_modified = False
        return modified_data

    def continue_breakpoint(self, modified_data=None):
        """继续执行断点"""
        if modified_data is not None:
            self.breakpoint_data = modified_data
            self.continue_with_modified = True
        self.breakpoint_event.set()

    def start(self):
        try:
            # 检查端口是否被占用
            if self.is_port_in_use(self.port):
                # 尝试释放端口
                if not self.force_release_port(self.port):
                    raise Exception(f"端口 {self.port} 已被占用且无法释放")
                # 等待端口释放
                time.sleep(1)

            self.running = True
            # 创建代理服务器套接字
            self.proxy_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.proxy_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                self.proxy_socket.bind((self.host, self.port))
            except socket.error as e:
                if e.errno == 10048:  # Windows系统端口被占用的错误码
                    raise Exception(f"端口 {self.port} 已被占用，请选择其他端口或关闭占用该端口的程序")
                raise e
            
            self.proxy_socket.listen(5)
            print(f"代理服务器启动在 {self.host}:{self.port}")

            # 启动接受连接的线程
            accept_thread = threading.Thread(target=self._accept_connections)
            accept_thread.daemon = True
            accept_thread.start()

        except Exception as e:
            self.running = False
            if self.proxy_socket:
                self.proxy_socket.close()
            error_msg = f"启动代理失败: {str(e)}"
            print(error_msg)
            messagebox.showerror("错误", error_msg)
            raise

    def _accept_connections(self):
        """接受客户端连接的方法"""
        while self.running:
            try:
                client_socket, address = self.proxy_socket.accept()
                print(f"接受来自 {address} 的连接")
                
                # 为每个客户端连接创建一个新线程
                client_thread = threading.Thread(
                    target=self._handle_client_connection,
                    args=(client_socket, address)
                )
                client_thread.daemon = True
                client_thread.start()
                
            except Exception as e:
                if self.running:
                    print(f"接受连接时出错: {e}")

    def _handle_client_connection(self, client_socket, address):
        """处理客户端连接的方法"""
        connection_id = f"{address[0]}:{address[1]}"
        self.connections[connection_id] = client_socket
        
        try:
            # 接收初始数据以确定是HTTP还是HTTPS
            initial_data = client_socket.recv(4096)
            if not initial_data:
                return

            # 安全检查
            is_safe, reason = self.security_manager.check_request(address[0], initial_data)
            if not is_safe:
                print(f"拒绝来自 {address[0]} 的请求: {reason}")
                # 发送403响应
                response = (
                    f"HTTP/1.1 403 Forbidden\r\n"
                    f"Content-Type: text/plain\r\n"
                    f"Connection: close\r\n"
                    f"\r\n"
                    f"Request blocked: {reason}"
                ).encode()
                client_socket.send(response)
                return

            # 检查是否是HTTPS CONNECT请求
            if initial_data.startswith(b'CONNECT'):
                # 解析主机名和端口
                first_line = initial_data.decode('utf-8').split('\r\n')[0]
                _, host_port, _ = first_line.split(' ')
                hostname, port = host_port.split(':')
                port = int(port)

                # 记录CONNECT请求
                if self.callback:
                    traffic_entry = {
                        'timestamp': datetime.datetime.now().isoformat(),
                        'client_address': f"{address[0]}:{address[1]}",
                        'type': 'request',
                        'data': {
                            'method': 'CONNECT',
                            'path': f"{hostname}:{port}",
                            'headers': {},
                            'body': ''
                        }
                    }
                    self.callback(traffic_entry)

                # 发送200 Connection Established
                client_socket.send(b'HTTP/1.1 200 Connection Established\r\n\r\n')

                # 使用SSL包装客户端socket
                try:
                    context = self.create_ssl_context(hostname)
                    client_socket = context.wrap_socket(client_socket, server_side=True)
                except Exception as e:
                    print(f"SSL包装失败: {e}")
                    return

                # 创建到目标服务器的SSL连接
                server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                server_socket = ssl.create_default_context().wrap_socket(
                    server_socket, server_hostname=hostname
                )
                server_socket.connect((hostname, port))
                
                # 保存服务器连接
                self.connections[f"{connection_id}_server"] = server_socket
            else:
                # 处理普通HTTP请求
                request_data = initial_data
                http_request = HTTPMessage(request_data)
                
                # 记录请求
                if self.callback:
                    traffic_entry = {
                        'timestamp': datetime.datetime.now().isoformat(),
                        'client_address': f"{address[0]}:{address[1]}",
                        'type': 'request',
                        'data': {
                            'method': http_request.method,
                            'path': http_request.path,
                            'headers': http_request.headers,
                            'body': http_request.body
                        }
                    }
                    self.callback(traffic_entry)

                # 解析目标服务器地址
                host = http_request.headers.get('Host', '').split(':')[0]
                port = int(http_request.headers.get('Host', '').split(':')[1]) if ':' in http_request.headers.get('Host', '') else 80

                # 连接到目标服务器
                server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                server_socket.connect((host, port))
                server_socket.send(request_data)

                # 接收服务器响应
                response_data = b""
                while True:
                    try:
                        chunk = server_socket.recv(4096)
                        if not chunk:
                            break
                        response_data += chunk
                    except socket.timeout:
                        break
                    except Exception as e:
                        print(f"接收响应数据时出错: {e}")
                        break

                if response_data:
                    # 解析HTTP响应
                    http_response = HTTPMessage(response_data)
                    
                    # 记录响应
                    if self.callback:
                        traffic_entry = {
                            'timestamp': datetime.datetime.now().isoformat(),
                            'server_address': f"{host}:{port}",
                            'type': 'response',
                            'data': {
                                'status_code': http_response.status_code,
                                'status_text': http_response.status_text,
                                'headers': http_response.headers,
                                'body': http_response.body
                            }
                        }
                        self.callback(traffic_entry)

                    # 发送响应到客户端
                    client_socket.send(response_data)
                
                # 保存服务器连接
                self.connections[f"{connection_id}_server"] = server_socket

            # 创建双向数据转发线程
            client_to_server = threading.Thread(
                target=self._forward_data,
                args=(client_socket, server_socket, f"{connection_id}_c2s"),
                daemon=True
            )
            server_to_client = threading.Thread(
                target=self._forward_data,
                args=(server_socket, client_socket, f"{connection_id}_s2c"),
                daemon=True
            )
            
            client_to_server.start()
            server_to_client.start()

        except Exception as e:
            print(f"处理连接时出错: {e}")
        finally:
            # 连接出错时清理
            if not self.running:
                self._cleanup_connection(connection_id)

    def _forward_data(self, source, destination, connection_id):
        """转发数据"""
        try:
            while self.running:
                try:
                    data = source.recv(4096)
                    if not data:
                        break

                    # 记录数据
                    if self.callback and len(data) > 0:
                        try:
                            http_message = HTTPMessage(data)
                            traffic_entry = {
                                'timestamp': datetime.datetime.now().isoformat(),
                                'type': 'request' if http_message.is_request else 'response',
                                'data': {
                                    'method': http_message.method,
                                    'path': http_message.path,
                                    'status_code': http_message.status_code,
                                    'status_text': http_message.status_text,
                                    'headers': http_message.headers,
                                    'body': http_message.body
                                }
                            }
                            self.callback(traffic_entry)
                        except Exception as e:
                            print(f"解析转发数据时出错: {e}")

                    destination.send(data)
                except (socket.timeout, ssl.SSLError) as e:
                    if isinstance(e, socket.timeout):
                        continue
                    else:
                        break
                except Exception as e:
                    print(f"转发数据时出错: {e}")
                    break
        finally:
            self._cleanup_connection(connection_id.rsplit('_', 1)[0])

    def _cleanup_connection(self, connection_id):
        """清理连接"""
        # 清理客户端连接
        if connection_id in self.connections:
            try:
                self.connections[connection_id].close()
            except:
                pass
            del self.connections[connection_id]
        
        # 清理服务器连接
        server_conn_id = f"{connection_id}_server"
        if server_conn_id in self.connections:
            try:
                self.connections[server_conn_id].close()
            except:
                pass
            del self.connections[server_conn_id]

    def stop(self):
        self.running = False
        if self.proxy_socket:
            try:
                # 关闭所有连接
                for conn in self.connections.values():
                    try:
                        conn.close()
                    except:
                        pass
                self.connections.clear()
                
                # 关闭代理服务器socket
                self.proxy_socket.close()
                self.proxy_socket = None
            except:
                pass
        print("代理服务器已停止")

class NetworkSnifferGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("网络抓包工具")
        self.root.geometry("1200x800")
        
        # 添加监听所有端口的选项
        self.listen_all_ports = tk.BooleanVar(value=False)
        
        # 添加自动滚动选项
        self.auto_scroll = tk.BooleanVar(value=True)
        
        # 添加断点控制
        self.break_on_request = tk.BooleanVar(value=False)
        self.break_on_response = tk.BooleanVar(value=False)
        
        # 过滤器变量
        self.filter_text = tk.StringVar()
        self.filter_method = tk.StringVar(value="ALL")
        self.filter_status = tk.StringVar(value="ALL")
        
        # 存储所有请求的列表
        self.all_requests = []
        self.filtered_requests = []
        
        # 当前断点的请求/响应
        self.current_breakpoint = None
        
        # 设置主题样式
        style = ttk.Style()
        style.configure("Treeview", rowheight=25)
        style.configure("TButton", padding=5)
        
        self.sniffer = None
        self.traffic_queue = queue.Queue()
        self.setup_gui()
        self.check_traffic_queue()

        # 添加断点编辑窗口
        self.breakpoint_window = None
        self.breakpoint_text = None

    def setup_gui(self):
        # 主分割面板
        self.main_paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        self.main_paned.pack(fill=tk.BOTH, expand=True)

        # 左侧面板
        left_frame = ttk.Frame(self.main_paned)
        self.main_paned.add(left_frame, weight=1)

        # 控制面板
        control_frame = ttk.LabelFrame(left_frame, text="控制面板", padding="5 5 5 5")
        control_frame.pack(fill=tk.X, padx=5, pady=5)

        # 网格布局控制面板
        ttk.Label(control_frame, text="主机:").grid(row=0, column=0, padx=5, pady=5)
        self.host_entry = ttk.Entry(control_frame, width=15)
        self.host_entry.insert(0, "127.0.0.1")
        self.host_entry.grid(row=0, column=1, padx=5, pady=5)

        ttk.Label(control_frame, text="端口:").grid(row=0, column=2, padx=5, pady=5)
        self.port_entry = ttk.Entry(control_frame, width=6)
        self.port_entry.insert(0, "8888")
        self.port_entry.grid(row=0, column=3, padx=5, pady=5)

        # 添加过滤器框架
        filter_frame = ttk.LabelFrame(left_frame, text="过滤器", padding="5 5 5 5")
        filter_frame.pack(fill=tk.X, padx=5, pady=5)

        ttk.Label(filter_frame, text="URL:").grid(row=0, column=0, padx=5, pady=5)
        self.filter_entry = ttk.Entry(filter_frame, textvariable=self.filter_text)
        self.filter_entry.grid(row=0, column=1, padx=5, pady=5, sticky='ew')
        self.filter_entry.bind('<Return>', lambda e: self.apply_filters())

        ttk.Label(filter_frame, text="方法:").grid(row=0, column=2, padx=5, pady=5)
        methods = ['ALL', 'GET', 'POST', 'PUT', 'DELETE']
        self.method_combo = ttk.Combobox(filter_frame, values=methods, textvariable=self.filter_method, width=8)
        self.method_combo.grid(row=0, column=3, padx=5, pady=5)
        self.method_combo.bind('<<ComboboxSelected>>', lambda e: self.apply_filters())

        ttk.Label(filter_frame, text="状态:").grid(row=0, column=4, padx=5, pady=5)
        status = ['ALL', '2xx', '3xx', '4xx', '5xx']
        self.status_combo = ttk.Combobox(filter_frame, values=status, textvariable=self.filter_status, width=8)
        self.status_combo.grid(row=0, column=5, padx=5, pady=5)
        self.status_combo.bind('<<ComboboxSelected>>', lambda e: self.apply_filters())

        # 添加断点控制框架
        breakpoint_frame = ttk.LabelFrame(left_frame, text="断点控制", padding="5 5 5 5")
        breakpoint_frame.pack(fill=tk.X, padx=5, pady=5)

        ttk.Checkbutton(breakpoint_frame, text="请求断点", variable=self.break_on_request).pack(side=tk.LEFT, padx=5)
        ttk.Checkbutton(breakpoint_frame, text="响应断点", variable=self.break_on_response).pack(side=tk.LEFT, padx=5)
        ttk.Checkbutton(breakpoint_frame, text="自动滚动", variable=self.auto_scroll).pack(side=tk.LEFT, padx=5)

        button_frame = ttk.Frame(breakpoint_frame)
        button_frame.pack(side=tk.RIGHT, padx=5)

        self.start_button = ttk.Button(button_frame, text="开始监听", command=self.toggle_monitoring)
        self.start_button.pack(side=tk.LEFT, padx=5)

        self.clear_button = ttk.Button(button_frame, text="清除记录", command=self.clear_display)
        self.clear_button.pack(side=tk.LEFT, padx=5)

        self.save_button = ttk.Button(button_frame, text="保存会话", command=self.save_session)
        self.save_button.pack(side=tk.LEFT, padx=5)

        self.load_button = ttk.Button(button_frame, text="加载会话", command=self.load_session)
        self.load_button.pack(side=tk.LEFT, padx=5)

        # 在控制面板添加监听所有端口的复选框
        self.listen_all_checkbox = ttk.Checkbutton(
            control_frame, 
            text="监听所有端口",
            variable=self.listen_all_ports,
            command=self.toggle_port_entry
        )
        self.listen_all_checkbox.grid(row=1, column=0, columnspan=2, padx=5, pady=5)

        # 请求列表
        list_frame = ttk.LabelFrame(left_frame, text="请求列表", padding="5 5 5 5")
        list_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # 创建树形视图和滚动条
        self.traffic_tree = ttk.Treeview(
            list_frame, 
            columns=("时间", "类型", "方法/状态", "路径/主机", "状态码"),
            show="headings"
        )
        vsb = ttk.Scrollbar(list_frame, orient="vertical", command=self.traffic_tree.yview)
        hsb = ttk.Scrollbar(list_frame, orient="horizontal", command=self.traffic_tree.xview)
        self.traffic_tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set)

        # 网格布局树形视图和滚动条
        self.traffic_tree.grid(column=0, row=0, sticky='nsew')
        vsb.grid(column=1, row=0, sticky='ns')
        hsb.grid(column=0, row=1, sticky='ew')

        list_frame.grid_columnconfigure(0, weight=1)
        list_frame.grid_rowconfigure(0, weight=1)

        # 设置列
        self.traffic_tree.heading("时间", text="时间")
        self.traffic_tree.heading("类型", text="类型")
        self.traffic_tree.heading("方法/状态", text="方法/状态")
        self.traffic_tree.heading("路径/主机", text="路径/主机")
        self.traffic_tree.heading("状态码", text="状态码")
        
        self.traffic_tree.column("时间", width=100)
        self.traffic_tree.column("类型", width=60)
        self.traffic_tree.column("方法/状态", width=80)
        self.traffic_tree.column("路径/主机", width=300)
        self.traffic_tree.column("状态码", width=60)

        # 绑定选择事件
        self.traffic_tree.bind('<<TreeviewSelect>>', self.on_select_request)

        # 右侧详细信息面板
        detail_frame = ttk.LabelFrame(self.main_paned, text="请求详情", padding="5 5 5 5")
        self.main_paned.add(detail_frame, weight=1)

        # 使用Notebook创建选项卡
        self.detail_notebook = ttk.Notebook(detail_frame)
        self.detail_notebook.pack(fill=tk.BOTH, expand=True)

        # 请求头选项卡
        self.headers_text = scrolledtext.ScrolledText(self.detail_notebook, wrap=tk.WORD)
        self.detail_notebook.add(self.headers_text, text='请求头')

        # 请求体选项卡
        self.body_text = scrolledtext.ScrolledText(self.detail_notebook, wrap=tk.WORD)
        self.detail_notebook.add(self.body_text, text='请求/响应体')

        # 原始数据选项卡
        self.raw_text = scrolledtext.ScrolledText(self.detail_notebook, wrap=tk.WORD)
        self.detail_notebook.add(self.raw_text, text='原始数据')

        # 添加安全设置框架
        security_frame = ttk.LabelFrame(left_frame, text="安全设置", padding="5 5 5 5")
        security_frame.pack(fill=tk.X, padx=5, pady=5)

        # 安全设置按钮
        ttk.Button(
            security_frame,
            text="安全配置",
            command=self.show_security_settings
        ).pack(side=tk.LEFT, padx=5)

        ttk.Button(
            security_frame,
            text="查看黑名单",
            command=self.show_blacklist
        ).pack(side=tk.LEFT, padx=5)

        # 添加安全日志显示
        self.security_log = scrolledtext.ScrolledText(
            security_frame,
            height=3,
            wrap=tk.WORD
        )
        self.security_log.pack(fill=tk.X, padx=5, pady=5)

    def toggle_port_entry(self):
        if self.listen_all_ports.get():
            self.port_entry.configure(state='disabled')
        else:
            self.port_entry.configure(state='normal')

    def toggle_monitoring(self):
        try:
            if self.start_button["text"] == "开始监听":
                host = self.host_entry.get()
                port = 0 if self.listen_all_ports.get() else int(self.port_entry.get())
                self.sniffer = NetworkSniffer(
                    host=host, 
                    port=port, 
                    callback=self.update_traffic_display,
                    listen_all=self.listen_all_ports.get()
                )
                self.sniffer.start()
                self.start_button["text"] = "停止监听"
                port_info = "所有端口" if self.listen_all_ports.get() else f"端口 {port}"
                messagebox.showinfo("成功", f"开始监听 {host}:{port_info}")

                # 设置断点状态
                self.break_on_request.trace_add('write', self._update_breakpoint_state)
                self.break_on_response.trace_add('write', self._update_breakpoint_state)
            else:
                if self.sniffer:
                    self.sniffer.stop()
                self.start_button["text"] = "开始监听"
                messagebox.showinfo("成功", "停止监听")
        except Exception as e:
            messagebox.showerror("错误", f"操作失败: {str(e)}")

    def _update_breakpoint_state(self, *args):
        """更新断点状态"""
        if self.sniffer:
            self.sniffer.set_breakpoint(
                request=self.break_on_request.get(),
                response=self.break_on_response.get()
            )

    def show_breakpoint_editor(self, data):
        """显示断点编辑窗口"""
        if self.breakpoint_window:
            self.breakpoint_window.destroy()

        self.breakpoint_window = tk.Toplevel(self.root)
        self.breakpoint_window.title("断点编辑器")
        self.breakpoint_window.geometry("800x600")

        # 创建编辑区
        self.breakpoint_text = scrolledtext.ScrolledText(self.breakpoint_window, wrap=tk.WORD)
        self.breakpoint_text.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # 解析并格式化数据
        try:
            message = HTTPMessage(data)
            formatted_text = []
            
            if message.is_request:
                formatted_text.append(f"{message.method} {message.path} HTTP/1.1")
            else:
                formatted_text.append(f"HTTP/1.1 {message.status_code} {message.status_text}")
            
            for key, value in message.headers.items():
                formatted_text.append(f"{key}: {value}")
            
            formatted_text.append("")
            formatted_text.append(message.body)
            
            self.breakpoint_text.insert('1.0', '\n'.join(formatted_text))
        except:
            # 如果解析失败，直接显示原始数据
            try:
                self.breakpoint_text.insert('1.0', data.decode('utf-8', errors='ignore'))
            except:
                self.breakpoint_text.insert('1.0', str(data))

        # 添加按钮
        button_frame = ttk.Frame(self.breakpoint_window)
        button_frame.pack(fill=tk.X, padx=5, pady=5)

        ttk.Button(
            button_frame, 
            text="继续执行", 
            command=lambda: self.continue_breakpoint()
        ).pack(side=tk.LEFT, padx=5)

        ttk.Button(
            button_frame, 
            text="修改并继续", 
            command=lambda: self.continue_breakpoint(modified=True)
        ).pack(side=tk.LEFT, padx=5)

    def continue_breakpoint(self, modified=False):
        """继续执行断点"""
        if not self.sniffer:
            return

        if modified and self.breakpoint_text:
            # 获取修改后的数据
            modified_data = self.breakpoint_text.get('1.0', tk.END).strip()
            # 转换为字节
            modified_data = modified_data.encode('utf-8')
            self.sniffer.continue_breakpoint(modified_data)
        else:
            self.sniffer.continue_breakpoint()

        if self.breakpoint_window:
            self.breakpoint_window.destroy()
            self.breakpoint_window = None
            self.breakpoint_text = None

    def update_traffic_display(self, traffic_entry):
        """更新流量显示"""
        # 记录安全事件
        if 'security_event' in traffic_entry:
            self.log_security_event(traffic_entry['security_event'])

        # 检查是否需要显示断点编辑器
        if self.sniffer and (
            (traffic_entry['type'] == 'request' and self.break_on_request.get()) or
            (traffic_entry['type'] == 'response' and self.break_on_response.get())
        ):
            # 重建原始数据
            if traffic_entry['type'] == 'request':
                data = (
                    f"{traffic_entry['data']['method']} {traffic_entry['data']['path']} HTTP/1.1\r\n" +
                    '\r\n'.join(f"{k}: {v}" for k, v in traffic_entry['data']['headers'].items()) +
                    "\r\n\r\n" +
                    traffic_entry['data']['body']
                ).encode('utf-8')
            else:
                data = (
                    f"HTTP/1.1 {traffic_entry['data']['status_code']} {traffic_entry['data']['status_text']}\r\n" +
                    '\r\n'.join(f"{k}: {v}" for k, v in traffic_entry['data']['headers'].items()) +
                    "\r\n\r\n" +
                    traffic_entry['data']['body']
                ).encode('utf-8')
            
            self.root.after(0, lambda: self.show_breakpoint_editor(data))

        self.traffic_queue.put(traffic_entry)

    def check_traffic_queue(self):
        while not self.traffic_queue.empty():
            traffic_entry = self.traffic_queue.get()
            self.add_traffic_entry(traffic_entry)
        self.root.after(100, self.check_traffic_queue)

    def add_traffic_entry(self, traffic_entry, refresh=True):
        """添加流量条目到列表"""
        # 添加到完整请求列表
        if refresh:
            self.all_requests.append(traffic_entry)

        timestamp = traffic_entry['timestamp'].split('T')[1][:8]
        entry_type = traffic_entry['type']
        data = traffic_entry['data']
        
        if entry_type == 'request':
            method_status = data.get('method', '')
            path = data.get('path', '')
            status_code = ''
        else:
            method_status = data.get('status_text', '')
            path = data.get('path', '')
            status_code = data.get('status_code', '')

        item_id = self.traffic_tree.insert("", tk.END, values=(
            timestamp, entry_type, method_status, path, status_code
        ))
        
        # 只在需要时自动滚动
        if self.auto_scroll.get():
            self.traffic_tree.see(item_id)
            
        self.traffic_tree.item(item_id, tags=(json.dumps(traffic_entry),))

    def on_select_request(self, event):
        selected_items = self.traffic_tree.selection()
        if not selected_items:
            return

        item = selected_items[0]
        traffic_entry = json.loads(self.traffic_tree.item(item)['tags'][0])
        
        # 更新请求头显示
        self.headers_text.delete(1.0, tk.END)
        self.headers_text.insert(tk.END, f"{traffic_entry['type'].upper()} 头信息:\n\n")
        for key, value in traffic_entry['data'].get('headers', {}).items():
            self.headers_text.insert(tk.END, f"{key}: {value}\n")
            
        # 更新请求/响应体显示
        self.body_text.delete(1.0, tk.END)
        if traffic_entry['type'] == 'request':
            self.body_text.insert(tk.END, "请求体:\n\n")
        else:
            self.body_text.insert(tk.END, "响应体:\n\n")
        
        body_content = traffic_entry['data'].get('body', '')
        try:
            # 尝试格式化JSON
            json_data = json.loads(body_content)
            body_content = json.dumps(json_data, indent=2, ensure_ascii=False)
        except:
            pass
            
        self.body_text.insert(tk.END, body_content)

        # 更新原始数据显示
        self.raw_text.delete(1.0, tk.END)
        self.raw_text.insert(tk.END, f"完整 {traffic_entry['type'].upper()} 数据:\n\n")
        if traffic_entry['type'] == 'request':
            self.raw_text.insert(tk.END, f"{traffic_entry['data'].get('method', '')} {traffic_entry['data'].get('path', '')} HTTP/1.1\n")
        else:
            self.raw_text.insert(tk.END, f"HTTP/1.1 {traffic_entry['data'].get('status_code', '')} {traffic_entry['data'].get('status_text', '')}\n")
        
        for key, value in traffic_entry['data'].get('headers', {}).items():
            self.raw_text.insert(tk.END, f"{key}: {value}\n")
        self.raw_text.insert(tk.END, f"\n{body_content}")

    def clear_display(self):
        self.traffic_tree.delete(*self.traffic_tree.get_children())
        self.headers_text.delete(1.0, tk.END)
        self.body_text.delete(1.0, tk.END)
        self.raw_text.delete(1.0, tk.END)

    def apply_filters(self):
        """应用过滤器到请求列表"""
        self.filtered_requests = []
        filter_text = self.filter_text.get().lower()
        filter_method = self.filter_method.get()
        filter_status = self.filter_status.get()

        for request in self.all_requests:
            # 检查URL过滤器
            if filter_text and filter_text not in request.get('data', {}).get('path', '').lower():
                continue

            # 检查方法过滤器
            if filter_method != 'ALL' and request.get('data', {}).get('method') != filter_method:
                continue

            # 检查状态码过滤器
            if filter_status != 'ALL':
                status_code = request.get('data', {}).get('status_code', '')
                if not status_code.startswith(filter_status[0]):
                    continue

            self.filtered_requests.append(request)

        self.refresh_traffic_tree()

    def refresh_traffic_tree(self):
        """刷新请求列表显示"""
        self.traffic_tree.delete(*self.traffic_tree.get_children())
        for traffic_entry in self.filtered_requests:
            self.add_traffic_entry(traffic_entry, refresh=False)

    def save_session(self):
        """保存当前会话到文件"""
        file_path = filedialog.asksaveasfilename(
            defaultextension=".json",
            filetypes=[("JSON files", "*.json"), ("All files", "*.*")]
        )
        if file_path:
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(self.all_requests, f, ensure_ascii=False, indent=2)
                messagebox.showinfo("成功", "会话已保存")
            except Exception as e:
                messagebox.showerror("错误", f"保存会话失败: {str(e)}")

    def load_session(self):
        """从文件加载会话"""
        file_path = filedialog.askopenfilename(
            filetypes=[("JSON files", "*.json"), ("All files", "*.*")]
        )
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    self.all_requests = json.load(f)
                self.apply_filters()
                messagebox.showinfo("成功", "会话已加载")
            except Exception as e:
                messagebox.showerror("错误", f"加载会话失败: {str(e)}")

    def replay_request(self):
        """重放选中的请求"""
        selected_items = self.traffic_tree.selection()
        if not selected_items:
            return

        item = selected_items[0]
        traffic_entry = json.loads(self.traffic_tree.item(item)['tags'][0])
        if traffic_entry['type'] != 'request':
            messagebox.showwarning("警告", "只能重放请求")
            return

        try:
            # 创建新的请求
            host = traffic_entry['data']['headers'].get('Host', '').split(':')[0]
            port = int(traffic_entry['data']['headers'].get('Host', '').split(':')[1]) if ':' in traffic_entry['data']['headers'].get('Host', '') else 80

            # 连接到目标服务器
            server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            server_socket.settimeout(10)
            server_socket.connect((host, port))

            # 重建请求数据
            request_lines = [
                f"{traffic_entry['data']['method']} {traffic_entry['data']['path']} HTTP/1.1"
            ]
            for key, value in traffic_entry['data']['headers'].items():
                request_lines.append(f"{key}: {value}")
            request_lines.append("")
            request_lines.append(traffic_entry['data']['body'])
            request_data = '\r\n'.join(request_lines).encode()

            # 发送请求
            server_socket.send(request_data)
            messagebox.showinfo("成功", "请求已重放")
            server_socket.close()
        except Exception as e:
            messagebox.showerror("错误", f"重放请求失败: {str(e)}")

    def show_security_settings(self):
        """显示安全设置窗口"""
        if not self.sniffer:
            messagebox.showwarning("警告", "请先启动监听")
            return

        settings_window = tk.Toplevel(self.root)
        settings_window.title("安全配置")
        settings_window.geometry("400x500")

        # 创建配置输入框
        frame = ttk.Frame(settings_window, padding="10")
        frame.pack(fill=tk.BOTH, expand=True)

        row = 0
        entries = {}
        for key, value in self.sniffer.security_manager.config.items():
            ttk.Label(frame, text=key).grid(row=row, column=0, padx=5, pady=5)
            entry = ttk.Entry(frame)
            entry.insert(0, str(value))
            entry.grid(row=row, column=1, padx=5, pady=5)
            entries[key] = entry
            row += 1

        # 添加攻击特征模式编辑区
        ttk.Label(frame, text="攻击特征模式:").grid(row=row, column=0, columnspan=2, padx=5, pady=5)
        row += 1
        
        patterns_text = scrolledtext.ScrolledText(frame, height=10)
        patterns_text.grid(row=row, column=0, columnspan=2, padx=5, pady=5)
        patterns_text.insert('1.0', '\n'.join(self.sniffer.security_manager.attack_patterns))

        def save_settings():
            try:
                # 更新配置
                for key, entry in entries.items():
                    value = entry.get()
                    if key in ['max_requests_per_minute', 'max_payload_size', 'max_uri_length', 'max_header_count', 'max_header_size']:
                        self.sniffer.security_manager.config[key] = int(value)
                    else:
                        self.sniffer.security_manager.config[key] = float(value)

                # 更新攻击特征
                patterns = patterns_text.get('1.0', tk.END).strip().split('\n')
                self.sniffer.security_manager.attack_patterns = [p for p in patterns if p]

                messagebox.showinfo("成功", "安全配置已更新")
                settings_window.destroy()
            except Exception as e:
                messagebox.showerror("错误", f"保存配置失败: {str(e)}")

        ttk.Button(
            frame,
            text="保存",
            command=save_settings
        ).grid(row=row+1, column=0, columnspan=2, pady=10)

    def show_blacklist(self):
        """显示黑名单窗口"""
        if not self.sniffer:
            messagebox.showwarning("警告", "请先启动监听")
            return

        blacklist_window = tk.Toplevel(self.root)
        blacklist_window.title("IP黑名单")
        blacklist_window.geometry("300x400")

        frame = ttk.Frame(blacklist_window, padding="10")
        frame.pack(fill=tk.BOTH, expand=True)

        # 创建列表框
        listbox = tk.Listbox(frame)
        listbox.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # 添加黑名单IP
        for ip in sorted(self.sniffer.security_manager.blacklist):
            listbox.insert(tk.END, ip)

        def remove_selected():
            selection = listbox.curselection()
            if not selection:
                return
            ip = listbox.get(selection[0])
            self.sniffer.security_manager.remove_from_blacklist(ip)
            listbox.delete(selection[0])

        ttk.Button(
            frame,
            text="移除所选IP",
            command=remove_selected
        ).pack(pady=5)

    def log_security_event(self, message: str):
        """记录安全事件"""
        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
        self.security_log.insert(tk.END, f"[{timestamp}] {message}\n")
        self.security_log.see(tk.END)

def main():
    try:
        root = tk.Tk()
        app = NetworkSnifferGUI(root)
        root.mainloop()
    except Exception as e:
        messagebox.showerror("错误", f"程序启动失败: {str(e)}")

if __name__ == "__main__":
    main()
